#!/usr/bin/env python3
"""
磁盘存储监控模块
监控磁盘空间使用情况，提供存储策略建议
"""

import os
import shutil
import psutil
import time
from typing import Dict, Any, Optional, List
from dataclasses import dataclass
from enum import Enum

class StorageLevel(Enum):
    """存储空间级别"""
    CRITICAL = "critical"      # 严重不足 (< 5%)
    LOW = "low"               # 空间不足 (5-15%)
    NORMAL = "normal"         # 正常 (15-50%)
    HIGH = "high"             # 充足 (50-80%)
    EXCESS = "excess"         # 过剩 (> 80%)

@dataclass
class DiskInfo:
    """磁盘信息"""
    total: int           # 总空间 (字节)
    used: int            # 已使用空间 (字节)
    free: int            # 可用空间 (字节)
    percentage: float    # 使用百分比
    level: StorageLevel  # 存储级别

@dataclass
class StorageRecommendation:
    """存储建议"""
    level: StorageLevel
    message: str
    actions: List[str]
    buffer_size_factor: float  # 缓冲区大小调整因子
    page_cleanup_threshold: int  # 页面清理阈值

class DiskMonitor:
    """磁盘存储监控器"""
    
    def __init__(self, data_directory: str = "data"):
        self.data_directory = data_directory
        self.monitoring = False
        self.monitor_interval = 30  # 监控间隔(秒)
        self.history: List[DiskInfo] = []
        self.max_history = 100
        
        # 存储级别阈值
        self.thresholds = {
            StorageLevel.CRITICAL: 5.0,
            StorageLevel.LOW: 15.0,
            StorageLevel.NORMAL: 50.0,
            StorageLevel.HIGH: 80.0,
            StorageLevel.EXCESS: 100.0
        }
    
    def get_disk_info(self, path: str = None) -> DiskInfo:
        """获取磁盘信息"""
        if path is None:
            path = self.data_directory
        
        # 确保路径存在
        os.makedirs(path, exist_ok=True)
        
        # 获取磁盘使用情况
        total, used, free = shutil.disk_usage(path)
        percentage = (used / total) * 100
        
        # 确定存储级别
        level = self._determine_storage_level(percentage)
        
        disk_info = DiskInfo(
            total=total,
            used=used,
            free=free,
            percentage=percentage,
            level=level
        )
        
        # 添加到历史记录
        self.history.append(disk_info)
        if len(self.history) > self.max_history:
            self.history.pop(0)
        
        return disk_info
    
    def _determine_storage_level(self, percentage: float) -> StorageLevel:
        """确定存储级别"""
        if percentage >= self.thresholds[StorageLevel.EXCESS]:
            return StorageLevel.EXCESS
        elif percentage >= self.thresholds[StorageLevel.HIGH]:
            return StorageLevel.HIGH
        elif percentage >= self.thresholds[StorageLevel.NORMAL]:
            return StorageLevel.NORMAL
        elif percentage >= self.thresholds[StorageLevel.LOW]:
            return StorageLevel.LOW
        else:
            return StorageLevel.CRITICAL
    
    def get_storage_recommendation(self, disk_info: DiskInfo = None) -> StorageRecommendation:
        """获取存储建议"""
        if disk_info is None:
            disk_info = self.get_disk_info()
        
        if disk_info.level == StorageLevel.CRITICAL:
            return StorageRecommendation(
                level=StorageLevel.CRITICAL,
                message="磁盘空间严重不足！需要立即清理",
                actions=[
                    "清理临时文件",
                    "删除过期数据",
                    "压缩数据文件",
                    "减少缓冲区大小",
                    "启用紧急清理模式"
                ],
                buffer_size_factor=0.3,
                page_cleanup_threshold=50
            )
        
        elif disk_info.level == StorageLevel.LOW:
            return StorageRecommendation(
                level=StorageLevel.LOW,
                message="磁盘空间不足，建议进行清理",
                actions=[
                    "清理临时文件",
                    "删除过期数据",
                    "减少缓冲区大小",
                    "启用定期清理"
                ],
                buffer_size_factor=0.5,
                page_cleanup_threshold=100
            )
        
        elif disk_info.level == StorageLevel.NORMAL:
            return StorageRecommendation(
                level=StorageLevel.NORMAL,
                message="磁盘空间正常",
                actions=[
                    "保持当前设置",
                    "定期监控"
                ],
                buffer_size_factor=1.0,
                page_cleanup_threshold=200
            )
        
        elif disk_info.level == StorageLevel.HIGH:
            return StorageRecommendation(
                level=StorageLevel.HIGH,
                message="磁盘空间充足",
                actions=[
                    "可以增加缓冲区大小",
                    "减少清理频率"
                ],
                buffer_size_factor=1.2,
                page_cleanup_threshold=300
            )
        
        else:  # EXCESS
            return StorageRecommendation(
                level=StorageLevel.EXCESS,
                message="磁盘空间过剩",
                actions=[
                    "可以大幅增加缓冲区大小",
                    "减少清理频率",
                    "启用数据压缩"
                ],
                buffer_size_factor=1.5,
                page_cleanup_threshold=500
            )
    
    def get_database_size(self) -> int:
        """获取数据库文件总大小"""
        total_size = 0
        if os.path.exists(self.data_directory):
            for dirpath, dirnames, filenames in os.walk(self.data_directory):
                for filename in filenames:
                    filepath = os.path.join(dirpath, filename)
                    try:
                        total_size += os.path.getsize(filepath)
                    except OSError:
                        continue
        return total_size
    
    def get_largest_files(self, limit: int = 10) -> List[tuple]:
        """获取最大的文件列表"""
        files = []
        if os.path.exists(self.data_directory):
            for dirpath, dirnames, filenames in os.walk(self.data_directory):
                for filename in filenames:
                    filepath = os.path.join(dirpath, filename)
                    try:
                        size = os.path.getsize(filepath)
                        files.append((filepath, size))
                    except OSError:
                        continue
        
        # 按大小排序
        files.sort(key=lambda x: x[1], reverse=True)
        return files[:limit]
    
    def cleanup_temp_files(self) -> int:
        """清理临时文件"""
        cleaned_count = 0
        temp_extensions = ['.tmp', '.temp', '.log', '.bak']
        
        if os.path.exists(self.data_directory):
            for dirpath, dirnames, filenames in os.walk(self.data_directory):
                for filename in filenames:
                    if any(filename.endswith(ext) for ext in temp_extensions):
                        filepath = os.path.join(dirpath, filename)
                        try:
                            os.remove(filepath)
                            cleaned_count += 1
                        except OSError:
                            continue
        
        return cleaned_count
    
    def get_memory_usage(self) -> Dict[str, Any]:
        """获取内存使用情况"""
        memory = psutil.virtual_memory()
        return {
            'total': memory.total,
            'available': memory.available,
            'used': memory.used,
            'percentage': memory.percent,
            'free': memory.free
        }
    
    def format_size(self, size_bytes: int) -> str:
        """格式化文件大小"""
        if size_bytes == 0:
            return "0B"
        
        size_names = ["B", "KB", "MB", "GB", "TB"]
        i = 0
        while size_bytes >= 1024 and i < len(size_names) - 1:
            size_bytes /= 1024.0
            i += 1
        
        return f"{size_bytes:.2f}{size_names[i]}"
    
    def print_status(self):
        """打印存储状态"""
        disk_info = self.get_disk_info()
        recommendation = self.get_storage_recommendation(disk_info)
        db_size = self.get_database_size()
        memory_info = self.get_memory_usage()
        
        print("=" * 60)
        print("磁盘存储状态")
        print("=" * 60)
        print(f"数据目录: {self.data_directory}")
        print(f"磁盘总空间: {self.format_size(disk_info.total)}")
        print(f"磁盘已使用: {self.format_size(disk_info.used)} ({disk_info.percentage:.1f}%)")
        print(f"磁盘可用空间: {self.format_size(disk_info.free)}")
        print(f"存储级别: {disk_info.level.value}")
        print(f"数据库大小: {self.format_size(db_size)}")
        print()
        print("内存使用情况:")
        print(f"  总内存: {self.format_size(memory_info['total'])}")
        print(f"  已使用: {self.format_size(memory_info['used'])} ({memory_info['percentage']:.1f}%)")
        print(f"  可用内存: {self.format_size(memory_info['available'])}")
        print()
        print("存储建议:")
        print(f"  级别: {recommendation.level.value}")
        print(f"  消息: {recommendation.message}")
        print(f"  缓冲区调整因子: {recommendation.buffer_size_factor}")
        print(f"  页面清理阈值: {recommendation.page_cleanup_threshold}")
        print("  建议操作:")
        for action in recommendation.actions:
            print(f"    - {action}")
        
        # 显示最大的文件
        largest_files = self.get_largest_files(5)
        if largest_files:
            print()
            print("最大的文件:")
            for filepath, size in largest_files:
                relative_path = os.path.relpath(filepath, self.data_directory)
                print(f"  {relative_path}: {self.format_size(size)}")
    
    def start_monitoring(self):
        """开始监控"""
        self.monitoring = True
        print(f"开始磁盘监控，间隔: {self.monitor_interval}秒")
    
    def stop_monitoring(self):
        """停止监控"""
        self.monitoring = False
        print("停止磁盘监控")
    
    def get_history_summary(self) -> Dict[str, Any]:
        """获取历史记录摘要"""
        if not self.history:
            return {}
        
        levels = [info.level for info in self.history]
        percentages = [info.percentage for info in self.history]
        
        return {
            'total_samples': len(self.history),
            'current_level': levels[-1].value,
            'min_percentage': min(percentages),
            'max_percentage': max(percentages),
            'avg_percentage': sum(percentages) / len(percentages),
            'level_distribution': {
                level.value: levels.count(level) for level in StorageLevel
            }
        }